//==========================================================================
// Copyright (c) 2000-2008,  Elastos, Inc.  All Rights Reserved.
//==========================================================================

#include <mantle.h>
#include <stdlib.h>
#include "peimage.h"

// Interface functions
//
PInterface ModuleView::Probe(
    /* [in] */ REIID riid)
{
    if (riid == EIID_IInterface) {
        return (PInterface)(IModule *)this;
    }
    else if (riid == EIID_IModule) {
        return (IModule *)this;
    }
    else if (riid == EIID_CLASS_INFO) {
        return (PInterface)&ECLSID_CModule;
    }
    else if (riid == EIID_INTERFACE_INFO) {
        return (PInterface)&EIID_IModule;
    }

    return NULL;
}

UInt32 ModuleView::AddRef(void)
{
    assert(ModViewState_Available == m_modviewstate);
    return m_cRef.Increment();
}

UInt32 ModuleView::Release(void)
{
    Int32 lRef = m_cRef.Decrement();

    assert(lRef >= 0);

    if (0 == lRef) {
        assert(ModViewState_Available == m_modviewstate);
        m_modviewstate = ModViewState_Deleting;

        TryToSendReceiveDebugMsg(
            DEBUGGEE_UNLOAD_DLL, (UInt32)m_pModule, m_pProcess);

        m_pProcess->UnmapAndDeleteModView(this);
    }
    return (UInt32)lRef;
}

ECode ModuleView::Aggregate(
    /* [in] */ AggregateType type,
    /* [in] */ PInterface pObj)
{
    return E_NOT_IMPLEMENTED;
}

ECode ModuleView::GetDomain(
    /* [out] */ PInterface *ppObj)
{
    return E_NOT_IMPLEMENTED;
}

IInterface *ModuleView::GetObject()
{
    return this;
}

ECode ModuleView::GetBaseAddress(
    /* [out] */ Address * pBase)
{
    assert(m_pModule);

    if (NULL == pBase) return E_INVALID_ARGUMENT;

    *pBase = m_pModule->m_uImageBase;
    return NOERROR;
}

ECode ModuleView::GetEntryPointAddress(
    /* [out] */ Address * pDllMainStartup)
{
    assert(m_pModule);

    if (NULL == pDllMainStartup) return E_INVALID_ARGUMENT;

    *pDllMainStartup = m_pModule->m_uEntryPoint;
    return NOERROR;
}

ECode ModuleView::GetMemorySize(
    /* [out] */ MemorySize * pSize)
{
    assert(m_pModule);

    if (NULL == pSize) return E_INVALID_ARGUMENT;

    *pSize = m_pModule->m_uTotalSize;
    return NOERROR;
}

ECode ModuleView::GetName(
    /* [out] */ WStringBuf * pName)
{
    assert(m_pModule);

    if (pName == NULL || !pName->GetCapacity()) {
        return E_INVALID_ARGUMENT;
    }

    pName->Copy(m_pModule->m_wszName);
    return NOERROR;
}

ECode ModuleView::GetFullPath(
    /* [out] */ WStringBuf * pName)
{
    assert(m_pModule);

    if (pName == NULL || !pName->GetCapacity()) {
        return E_INVALID_ARGUMENT;
    }

    pName->Copy(m_pModule->m_wszFullPath);
    return NOERROR;
}

ECode ModuleView::GetProcAddressByName(
    /* [in] */ AString esName,
    /* [out] */ Address * pAddress)
{
    Address address;

    if (NULL == pAddress || NULL == esName) return E_INVALID_ARGUMENT;

    address = GetExportAddressByName(m_pModule, esName);
    if (NULL == address) {
        return E_DOES_NOT_EXIST;
    }
    else {
        *pAddress = address;
        return NOERROR;
    }
}

ECode ModuleView::GetProcAddressByOrdinal(
    /* [in] */ Int32 ordinal,
    /* [out] */ Address * pAddress)
{
    Address address;
    if (NULL == pAddress) return E_INVALID_ARGUMENT;

    address = GetExportAddressByOrdinal(m_pModule, ordinal);
    if (NULL == address) {
        return E_DOES_NOT_EXIST;
    }
    else {
        *pAddress = address;
        return NOERROR;
    }
}

ECode ModuleView::GetResourceSection(
    /* [out] */ Address * pAddress)
{
    Address address;

    assert(m_pModule);

    if (NULL == pAddress) return E_INVALID_ARGUMENT;

    if (NULL == m_pModule->m_pResource) return E_DOES_NOT_EXIST;
    address = m_pModule->m_uImageBase \
        + m_pModule->m_pResource->m_uVAddr;
    if (NULL == address) {
        return E_DOES_NOT_EXIST;
    }
    else {
        *pAddress = address;
        return NOERROR;
    }
}

ECode ModuleView::GetSectionByName(
    /* [in] */ AString esName,
    /* [out] */ Address * pAddress)
{
    UInt32 i;

    assert(m_pModule);
    assert(m_pModule->m_pSections);

    if (NULL == pAddress || NULL == esName) return E_INVALID_ARGUMENT;

    for (i = 0; i < m_pModule->m_cSections; i++) {
        if (!strcmp(m_pModule->m_pSections[i].m_szName, \
            esName)) {
            *pAddress = m_pModule->m_uImageBase \
                + m_pModule->m_pSections[i].m_uVAddr;
            return NOERROR;
        }
    }

    return E_DOES_NOT_EXIST;
}

PIMAGE_RESOURCE_DIRECTORY FindResDir(PIMAGE_RESOURCE_DIRECTORY prd,
    UInt32 dwResBase,
    wchar_t* pName)
{
    PIMAGE_RESOURCE_DIRECTORY_ENTRY prde;
    prde = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((UInt32)prd +
                sizeof(IMAGE_RESOURCE_DIRECTORY));

    if (*pName == L'#') {
        int id = _wtoi(pName + 1);
        prde += prd->uhNumberOfNamedEntries;
        for (int i = 0; i < prd->uhNumberOfIdEntries; i++) {
            if (prde->uhId == id) {
                if (prde->ulDataIsDirectory) {
                    return (PIMAGE_RESOURCE_DIRECTORY)
                             ((char*)dwResBase + prde->ulOffsetToDirectory);
                }
            }
            prde++;
        }
    }
    else {
        PIMAGE_RESOURCE_DIRECTORY_STRING prds = NULL;
        /* loop through all resource directory entry types */
        for (int i = 0; i < prd->uhNumberOfNamedEntries; i++) {
            if (prde->ulNameIsString) {
                prds = (PIMAGE_RESOURCE_DIRECTORY_STRING)
                                ((char*)dwResBase + prde->ulNameOffset);

                if (_wcsnicmp(prds->wszNameString, pName, prds->uhLength)
                        == 0) {
                    if (prde->ulDataIsDirectory) {
                        return (PIMAGE_RESOURCE_DIRECTORY)
                            ((char*)dwResBase + prde->ulOffsetToDirectory);
                    }
                }
            }
            prde++;
        }
    }

    return NULL;
}

PIMAGE_RESOURCE_DATA_ENTRY FindResEntry(PIMAGE_RESOURCE_DIRECTORY prd,
    UInt32 dwResBase)
{
    PIMAGE_RESOURCE_DIRECTORY_ENTRY prde;
    prde = (PIMAGE_RESOURCE_DIRECTORY_ENTRY)((UInt32)prd +
                sizeof(IMAGE_RESOURCE_DIRECTORY));

    if (!prde->ulDataIsDirectory) {
        return (PIMAGE_RESOURCE_DATA_ENTRY)
                        ((char*)dwResBase + prde->ulOffsetToData);
    }
    else {
        return NULL;
    }
}

ECode ModuleView::FindResource(
    /* [in] */ WString name,
    /* [in] */ WString type,
    /* [out] */ Address * pAddress,
    /* [out] */ MemorySize * pSize)
{
    PIMAGE_RESOURCE_DIRECTORY  prd, prd2, prd3;
    PIMAGE_RESOURCE_DATA_ENTRY pdata;
    UInt32 dwBase, dwResBase;

    ECode ec = NOERROR;

    if (NULL == pSize || NULL == pAddress
        || name.IsNull() || type.IsNull()) {
        return E_INVALID_ARGUMENT;
    }
    assert(m_pModule);
    dwBase = m_pModule->m_uImageBase;

    ec = GetResourceSection((Address *)&dwResBase);
    if (FAILED(ec)) {
        return ec;
    }

    prd = (PIMAGE_RESOURCE_DIRECTORY)dwResBase;
    /* set pointer to first resource type entry */
    prd2 = FindResDir(prd, dwResBase, (wchar_t*)(const wchar_t*)type);
    if (prd2 == NULL) return E_DOES_NOT_EXIST;

    prd3 = FindResDir(prd2, dwResBase, (wchar_t*)(const wchar_t*)name);
    if (prd3 == NULL) return E_DOES_NOT_EXIST;

    pdata = FindResEntry(prd3, dwResBase);
    if (pdata == NULL) return E_DOES_NOT_EXIST;

    *pAddress = (Address)((char*)dwBase + pdata->ulOffsetToData);
    *pSize = pdata->ulSize;

    return NOERROR;
}

// member functions
//
ModuleView::~ModuleView()
{
    if (m_pModule) m_pModule->Release();
}
